import Router from '@koa/router'
import { runtimeConfig } from '../../configs/runtimeConfig'
import { chatHandlerMap } from './chatHandler'
import { IReqData } from './interfaces'
import {
    getMessagesForCase,
    getMessagesForQuestion,
    getMessagesForQuestionPoint,
    getMessagesForWriting,
} from './getMessages'

// 案例问答controller
export const chatControllerForCase: Router.Middleware = async (ctx) => {
    const {
        user,
        context,
        knowledgeBase,
        enableKnowledgeBase,
        messages: reqMessages,
    } = ctx.request.body as IReqData
    if (!user) {
        ctx.status = 400
        ctx.body = 'message is required'
    }

    try {
        const messages = await getMessagesForCase({
            user,
            context,
            knowledgeBase,
            enableKnowledgeBase,
            messages: reqMessages,
        })

        const currentSDK = runtimeConfig.caseChatModel
            .sdk as keyof typeof chatHandlerMap

        const chatHandle = chatHandlerMap[currentSDK]

        const response = await chatHandle(messages, runtimeConfig.caseChatModel)
        ctx.set({ 'Content-Type': 'text/event-stream' })
        ctx.status = 200
        for await (const part of response) {
            ctx.res.write(
                `data: ${JSON.stringify({
                    text: part,
                })}\n\n`
            )
        }
    } catch (error) {
        console.log(error)
        ctx.status = 500
        ctx.body = { success: false, msg: '请求失败' }
    }
}

// 文章润色controller
export const chatControllerForRichText: Router.Middleware = async (ctx) => {
    const { user, context } = ctx.request.body as IReqData
    if (!user) {
        ctx.status = 400
        ctx.body = 'message is required'
    }

    ctx.set({ 'Content-Type': 'text/event-stream' })

    const messages = getMessagesForWriting({
        user,
        context,
    })

    const currentSDK = runtimeConfig.writingChatModel
        .sdk as keyof typeof chatHandlerMap

    const chatHandle = chatHandlerMap[currentSDK]

    const response = await chatHandle(messages, runtimeConfig.writingChatModel)
    ctx.status = 200
    // ctx.body = stream
    for await (const part of response) {
        // stream.write(`event: message\ndata: ${part.message.content}\n\n`)
        // stream.write(`data: ${part.message.content}\n\n`)
        ctx.res.write(
            `data: ${JSON.stringify({
                text: part,
            })}\n\n`
        )
    }
    // console.log('end')
    // stream.write('0')
    // 关闭connection
    // ctx.res.end()
}

// 知识点生成
export const chatControllerForQuestionPoint: Router.Middleware = async (
    ctx
) => {
    const {
        user,
        context,
        messages: reqMessages,
    } = ctx.request.body as IReqData
    if (!user) {
        ctx.status = 400
        ctx.body = 'message is required'
    }

    try {
        const messages = await getMessagesForQuestionPoint({
            user,
            context,
            messages: reqMessages,
        })

        const currentSDK = runtimeConfig.caseChatModel
            .sdk as keyof typeof chatHandlerMap

        const chatHandle = chatHandlerMap[currentSDK]

        const response = await chatHandle(messages, runtimeConfig.caseChatModel)
        ctx.set({ 'Content-Type': 'text/event-stream' })
        ctx.status = 200
        for await (const part of response) {
            ctx.res.write(
                `data: ${JSON.stringify({
                    text: part,
                })}\n\n`
            )
        }
    } catch (error) {
        console.log(error)
        ctx.status = 500
        ctx.body = { success: false, msg: '请求失败' }
    }
}

// 问题生成
export const chatControllerForQuestion: Router.Middleware = async (ctx) => {
    const {
        user,
        context,
        messages: reqMessages,
    } = ctx.request.body as IReqData
    if (!user) {
        ctx.status = 400
        ctx.body = 'message is required'
    }

    try {
        const messages = await getMessagesForQuestion({
            user,
            context,
            messages: reqMessages,
        })

        const currentSDK = runtimeConfig.caseChatModel
            .sdk as keyof typeof chatHandlerMap

        const chatHandle = chatHandlerMap[currentSDK]

        const response = await chatHandle(
            messages,
            runtimeConfig.caseChatModel,
            { outputJSON: true }
        )
        ctx.set({ 'Content-Type': 'text/event-stream' })
        ctx.status = 200
        for await (const part of response) {
            ctx.res.write(
                `data: ${JSON.stringify({
                    text: part,
                })}\n\n`
            )
        }
    } catch (error) {
        console.log(error)
        ctx.status = 500
        ctx.body = { success: false, msg: '请求失败' }
    }
}

// router.get('/sse', async (ctx) => {
//     try {
//         const stream = new PassThrough()
//         ctx.set({ 'Content-Type': 'text/event-stream' })
//         // 将歌词变成一个数组
//         let song = [
//             '我',
//             '懒',
//             '得',
//             '写',
//             '你',
//             '谷',
//             '搜',
//             '到',
//             '处',
//             '皆',
//             '只',
//             '因',
//             '你',
//             '太',
//             '美',
//             '浅',
//             '唱',
//             '动',
//             '人',
//             '说',
//             '不',
//             '出',
//             '我',
//             '试',
//             '着',
//             'end',
//         ]
//         let index = 0
//         setInterval(() => {
//             const str = song[index]
//             if (str) {
//                 // 方式1：格式是 data: xxx\n\n，事件要一对\n结束
//                 // stream.write(`data: ${song[index]}\n\n`)
//                 // 方式2,定义了event字段，自定义事件
//                 stream.write(`event: custom\ndata: ${song[index]}\n\n`)
//                 // 发送的数据为json字符串
//                 // stream.write(`event: custom\ndata: {"name": "${song[index]}"}\n\n`)
//             } else {
//                 stream.write('0')
//             }
//             index++
//         }, 500)

//         ctx.body = stream
//     } catch (error: any) {
//         error.status = error.status ? error.status : 500
//         ctx.throw(error.status, error)
//     }
// })

// router.get('/', async (ctx) => {
//     const stream = new PassThrough()
//     ctx.set({ 'Content-Type': 'text/event-stream' })
//     const response = await ollama.chat({
//         model: 'qwen2:7b',
//         messages: [{ role: 'user', content: '你好' }],
//         stream: true,
//     })
//     ctx.status = 200
//     ctx.body = stream
//     for await (const part of response) {
//         console.log(part.message.content)
//         // stream.write(`event: message\ndata: ${part.message.content}\n\n`)
//         stream.write(`data: ${part.message.content}\n\n`)
//     }
//     // console.log('end')
//     // stream.write('0')
//     // 关闭connection
//     // ctx.res.end()
// })
